"""
@author:王耀
@date:2021/9/18
"""

import cv2
import numpy as np
import math

def SobelKernel(size):
    """
    作用：一种通用的方法获取Sobel Kernal
    size(int)：Sobel kernal大小，size取奇数

    返回：x方向的kernel，y方向上的kernel
    """

    # 首先判断输入的sobel size是否合法
    # 再分别计算Pascal和变异Pascal三角形作为两个向量
    # 将这两个向量进行乘法运算得到两个矩阵，即为Sobel内核
    if size % 2 == 0:
        raise Exception('The kernel size of Sobel must be odd')
    def t(x):
        return math.factorial(x)
    def xy(m):
        A = []
        B = []
        for i in range(m - 1):
            n2 = int(t(m - 1) / (t(i) * t(m - i - 1)))
            n = int((t(m - 2)*(m - 2 * i - 1)) / (t(i) * t(m - i - 1)))
            A.append(n2)
            B.append(n)
        A.append(1)
        B.append(-1)
        return A, B

    sobelsmooth, sobeldiff = xy(size)
    sobelsmooth = np.array(sobelsmooth)
    sobeldiff = np.array(sobeldiff)
    return sobelsmooth.reshape(-1, 1) * sobeldiff, sobeldiff.reshape(-1, 1) * sobelsmooth



def SobelKernel3_5_7(size):
    """
    作用：一种简单的方法获取Sobel Kernal，用以加速
    size(int)：Sobel kernal大小，size取3、5、7

    返回：x方向的kernel，y方向上的kernel
    """

    # 判断size是否合法，直接手动构造两个向量，并相乘返回
    if size == 3:
        sobelsmooth = [1, 2, 1]
        sobeldiff = [1, 0, -1]
    elif size == 5:
        sobelsmooth = [1, 4, 6, 4, 1]
        sobeldiff = [1, 2, 0, -2, -1]
    elif size == 7:
        sobelsmooth = [1, 6, 15, 20, 15, 6, 1]
        sobeldiff = [1, 4, 5, 0, -5, -4, -1]

    else:
        raise Exception('Size ERROR in Sobel kernel, please use size = {3, 5, 7}.')


    sobelsmooth = np.array(sobelsmooth)
    sobeldiff = np.array(sobeldiff)
    return sobelsmooth.reshape(-1, 1) * sobeldiff, sobeldiff.reshape(-1, 1) * sobelsmooth



def Sobel_Kernel_four_direct(size = 3):
    """
    作用： 四个方向上的sobel滤波核
    :param size: 滤波核大小，限定为3

    :return:四个方向的Sobel滤波核
    """
    if size != 3:
        raise Exception('Sobel Size error')
    kernalx = np.array([[-1, 0, 1],
                       [-2, 0, 2],
                       [-1, 0, 1]])

    kernal45 = np.array([[0, 1, 2],
                       [-1, 0, 1],
                       [-2, -1, 0]])

    kernaly = np.array([[1, 2, 1],
                       [0, 0, 0],
                       [-1, -2, -1]])

    kernal135= np.array([[-2, -1, 0],
                         [-1, 0, 1],
                         [0, 1, 2]])

    return kernalx, kernaly, kernal45, kernal135



def SobelFilter(img, size = 3, directnum = 2):
    """
    作用：Sobel核心代码
    img：输入的图片
    size(int)：滤波kernal大小

    返回：（x、y方向分别的梯度也会返回，如果需要进行NMS的插值处理，则需要利用这个参数）图片像素梯度矩阵，图片像素梯度方向
    """
    global d45, d135, kernel135, kernel45
    if img is None:
        raise Exception('input image ERROR!')

    # sobel核的方向只有2或者4两种选择
    if directnum != 2 and directnum != 4:
        raise Exception('Please Check Your Sobel Direct...')

    if len(img.shape) == 2:
        H, W = img.shape
    else:
        raise Exception('Please input gray img image...')

    center = size // 2

    # 初始化x，y方向上的梯度矩阵
    dx = np.zeros((H + center * 2, W + center * 2), dtype=np.float)
    dx[center: center + H, center: center + W] = img.copy().astype(np.float)
    dy = np.zeros((H + center * 2, W + center * 2), dtype=np.float)
    dy[center: center + H, center: center + W] = img.copy().astype(np.float)

    # 如果方向设置为4，则初始化45°、135°方向上的梯度矩阵
    if directnum == 4:
        d45 = np.zeros((H + center * 2, W + center * 2), dtype=np.float)
        d45[center: center + H, center: center + W] = img.copy().astype(np.float)
        d135 = np.zeros((H + center * 2, W + center * 2), dtype=np.float)
        d135[center: center + H, center: center + W] = img.copy().astype(np.float)

    # 整体的梯度矩阵
    d = np.zeros((H + center * 2, W + center * 2), dtype=np.float)
    # 梯度方向矩阵
    direct = np.zeros((H + center * 2, W + center * 2), dtype=np.float)

    if directnum == 2:
        # kernelx, kernely = SobelKernel(size)  # 也可以使用该函数，获取更多大小的sobel kernel。
        kernelx, kernely = SobelKernel3_5_7(size)  # 获取两个方向的sobel核
    else:
        kernelx, kernely, kernel45, kernel135 = Sobel_Kernel_four_direct(size)  # 获取四个方向的sobel核，size限制为3
    
    restmp = dx.copy()

    # 逐一计算x、y方向上的梯度，并求出梯度和梯度方向
    # 分别存入d和drect中
    for h in range(H):
        for w in range(W):
            dx[center + h, center + w] = np.sum(kernelx * restmp[h: h + size, w: w + size])
            dy[center + h, center + w] = np.sum(kernely * restmp[h: h + size, w: w + size])
            if directnum == 4:
                d45[center + h, center + w] = np.sum(kernel45 * restmp[h: h + size, w: w + size])
                d135[center + h, center + w] = np.sum(kernel135 * restmp[h: h + size, w: w + size])

            # 计算整体的梯度
            if directnum == 2:
                d[center + h, center + w] = np.sqrt(np.square(dx[center + h, center + w]) + np.square(dy[center + h, center + w]))
            else:
                d[center + h, center + w] = np.sqrt(
                    np.square(dx[center + h, center + w]) + np.square(dy[center + h, center + w]) +
                    np.square(d45[center + h, center + w]) + np.square(d135[center + h, center + w])
                )

            # 计算梯度方向
            direct[center + h, center + w] = math.degrees(math.atan2(dy[center + h, center + w], dx[center + h, center + w]))
            # 负值统一变成正值
            if direct[center + h, center + w] < 0 : direct[center + h, center + w] += 360

    d = np.clip(d, 0, 255)
    d = d[center: center + H, center: center + W]
    dx = dx[center: center + H, center: center + W]
    dy = dy[center: center + H, center: center + W]
    direct = direct[center: center + H, center: center + W]

    print('执行完毕：{}方向sobel'.format(directnum))
    # 顺便返回dx、dy，以便进行NMS的插值处理
    return dx, dy, d, direct




#####################################
# if __name__ == '__main__':
#     print(SobelKernel3_5_7(3))
#     print(SobelKernel(9))